from keras.models import Model
from keras.losses import binary_crossentropy
from keras.metrics import binary_accuracy
from keras.callbacks import ModelCheckpoint, TensorBoard
from keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, Conv2DTranspose, concatenate
from keras.optimizers import Adam,RMSprop
import numpy as np
import glob
import cv2

# 定义 FCN 模型
def fcn_model(input_shape=(256, 256, 3), num_classes=3):
    inputs = Input(shape=input_shape)

    # 编码器
    conv1 = Conv2D(64, (3, 3), activation='relu', padding='same')(inputs)
    conv1 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv1)
    pool1 = MaxPooling2D((2, 2), strides=(2, 2))(conv1)

    conv2 = Conv2D(128, (3, 3), activation='relu', padding='same')(pool1)
    conv2 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv2)
    pool2 = MaxPooling2D((2, 2), strides=(2, 2))(conv2)

    conv3 = Conv2D(256, (3, 3), activation='relu', padding='same')(pool2)
    conv3 = Conv2D(256, (3, 3), activation='relu', padding='same')(conv3)
    pool3 = MaxPooling2D((2, 2), strides=(2, 2))(conv3)

    # 中间层
    conv4 = Conv2D(512, (3, 3), activation='relu', padding='same')(pool3)
    conv4 = Conv2D(512, (3, 3), activation='relu', padding='same')(conv4)
    pool4 = MaxPooling2D((2, 2), strides=(2, 2))(conv4)

    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same')(pool4)
    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same')(conv5)

    # 解码器
    up6 = Conv2D(256, (2, 2), activation='relu', padding='same')(UpSampling2D((2, 2))(conv5))
    merge6 = concatenate([conv4, up6], axis=-1)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same')(merge6)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same')(conv6)

    up7 = Conv2D(128, (2, 2), activation='relu', padding='same')(UpSampling2D((2, 2))(conv6))
    merge7 = concatenate([conv3, up7], axis=-1)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same')(merge7)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv7)

    up8 = Conv2D(64, (2, 2), activation='relu', padding='same')(UpSampling2D((2, 2))(conv7))
    merge8 = concatenate([conv2, up8], axis=-1)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same')(merge8)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv8)

    up9 = Conv2D(num_classes, (2, 2), activation='relu', padding='same')(UpSampling2D((2, 2))(conv8))
    conv9 = Conv2D(num_classes, (3, 3), activation='relu', padding='same')(up9)

    outputs = Conv2D(num_classes, (1, 1), activation='sigmoid' if num_classes==1 else 'softmax')(conv9)
    model = Model(inputs=inputs, outputs=outputs, name="fcn")
    return model

from keras.preprocessing.image import ImageDataGenerator
def generator(image_directory,target_size,batch_size,dict_data,seed):
    # 定义图像生成器
    datagen_image = ImageDataGenerator(**dict_data)
    datagen_label = ImageDataGenerator(**dict_data)
    data = zip(
        datagen_image.flow_from_directory(
            image_directory,
            classes=['image'],
            target_size=target_size,
            batch_size=batch_size,       #样本数不易过大
            seed=seed),
        datagen_label.flow_from_directory(
            image_directory,
            classes=["label"],
            color_mode='grayscale',
            target_size=target_size,
            batch_size=batch_size,       #样本数不易过大
            seed=seed)
        )
    for (img,label) in data:
        # cv2.imshow("image", img[0][0,:,:,:])
        # cv2.imshow("label", label[0][0,:,:,:])
        # cv2.waitKey(0)
        yield img[0],label[0]

if __name__ == "__main__":
    path = 'deep_net\\test'
    img = cv2.imread(glob.glob(path+"\\image\\*.jpg")[0])
    img = cv2.resize(img, (img.shape[1]//16*16, img.shape[0]//16*16))
    model = fcn_model(img.shape,1)
    if True:
        model.summary()
        data = generator(path,img.shape[:2], 5, dict(
            rescale=1./255,rotation_range=10,width_shift_range=0.2,height_shift_range=0.2,
            shear_range=0.2,zoom_range=0.2,horizontal_flip=True,fill_mode='constant'),1)
        model.compile(optimizer=RMSprop(learning_rate=0.001), loss=binary_crossentropy, metrics=['accuracy'])
        model.fit(data, steps_per_epoch=2, epochs=10, callbacks=
                    [
                        TensorBoard('deep_net/logs',write_images=True,histogram_freq=1,write_graph=False),
                        ModelCheckpoint("deep_net\\param\\fcn.keras"),
                    ])
    else:
        model.load_weights("deep_net\\param\\fcn.keras")
        for file in glob.glob(path+"\\image\\*.jpg")+glob.glob(path+"\\image\\*.png"):
            new_img = cv2.imread(file)/255.
            new_img = cv2.resize(new_img,(img.shape[1],img.shape[0]))
            new_img = new_img.reshape(1,*img.shape[:2],3)
            new_img = model.predict(new_img)
            cv2.imshow("",new_img[0,:,:,:])
            cv2.waitKey(1)